/* benchmark guru interfaces */
#include <Utilities/RandomGenerator.h>
#include <Configuration.h>
#include <mpi/communicator.h>
#include <mpi/datatype.h>
#include <iostream>
#include <fstream>
#include <Utilities/Timer.h>
#include <Message/OpenMP.h>
#include <OhmmsPETE/OhmmsMatrix.h>
#include <OhmmsData/hdf_datatype.h>
using namespace std;
using namespace APPNAMESPACE;


template<int D, typename ST=hsize_t>
struct DataView
{
  mpi::communicator& mycomm;
  ST count[D];
  ST offset[D];
  ST gcount[D];
  ST stride[D];
  size_t ndata;

  DataView(mpi::communicator& comm)
    : mycomm(comm), ndata(0)
    {
      for(int i=0; i<D; ++i) offset[i]=0;
      for(int i=0; i<D; ++i) stride[i]=1;
    }

  inline size_t size() const { return ndata;}

  inline void resize_last_dim()
  {
    count[D-1]*=2;
    gcount[D-1]*=2;
    ndata*=2;
  }

  //1d partition 
  void init(int n, bool iscomplex=false)
  {
    count[0]=n/mycomm.size();
    offset[0]=count[0]*mycomm.rank();
    ndata=count[0];
    for(int i=1; i<D; ++i) {count[i]=n; ndata*=n;}
    for(int i=0; i<D; ++i) gcount[i]=n;
    if(iscomplex) resize_last_dim();
  }

  //2d partition 
  void init(int n, int px, int py, bool iscomplex=false)
  {
    count[0]=n/px; 
    count[1]=n/py;
    ndata=count[0]*count[1];
    offset[0]=count[0]*(mycomm.rank()/py);
    offset[1]=count[1]*(mycomm.rank()%py);
    for(int i=2; i<D; ++i) {count[i]=n; ndata*=n;}
    for(int i=0; i<D; ++i) gcount[i]=n;
    if(iscomplex) resize_last_dim();
  }

  void print(ostream& os)
  {
    os << "rank " << mycomm.rank(); 
    for(int i=0; i<D; ++i) os << " " <<  offset[i];
    for(int i=0; i<D; ++i) os << " " <<  count[i];
    os << endl;
  }
};


template<typename T>
inline void write_pencils(size_t nitems, size_t nc, T* first, const string& rname,
    mpi::communicator& comm)
{
  char fname[128];
  int ierr;
  MPI_Offset offset=0;
  MPI_Info info;
  MPI_Status status;
  MPI_Datatype etype=mpi::get_mpi_datatype(*first);
  for(int i=0; i<nc; ++i)
  {
    sprintf(fname,"%s.%d",rname.c_str(),i);
    MPI_File fout;
    ierr=MPI_File_open(comm,fname,MPI_MODE_WRONLY|MPI_MODE_CREATE,MPI_INFO_NULL,&fout);
    if(ierr)
    {
      cerr << " Failed to open a mpi file " << fname << endl;
    }
    ierr=MPI_File_set_view(fout,offset,etype,etype,"native",info);
    ierr=MPI_File_write_at_all(fout,offset,(void*)(first+i*nitems),nitems,etype,&status);
    ierr=MPI_File_close(&fout);
  }
}

template<int D, typename T>
inline void write_pencils_h5(DataView<D>& view, T* first, int nc, const string& rname
    , bool buffered, bool collective=true)
{
  char fname[128];
  sprintf(fname,"%s.h5",rname.c_str());
  const hid_t etype=get_h5_datatype(*first);
  MPI_Info info=MPI_INFO_NULL;
  hid_t h_plist = H5Pcreate(H5P_FILE_ACCESS);
  H5Pset_fapl_mpio(h_plist,view.mycomm,info);
  hid_t h_file = H5Fcreate(fname,H5F_ACC_TRUNC,H5P_DEFAULT,h_plist);

  hid_t xfer_plist = H5Pcreate(H5P_DATASET_XFER);
  //Select either collective or indepedent
  if(collective)
    H5Pset_dxpl_mpio(xfer_plist,H5FD_MPIO_COLLECTIVE);
  else
    H5Pset_dxpl_mpio(xfer_plist,H5FD_MPIO_INDEPENDENT);
 
  //hsize_t buffsize=view.size()/view.count[0];
  hsize_t buffsize=view.size();
  if(buffered)
    hid_t b=H5Pset_buffer(xfer_plist, buffsize, NULL, NULL);

  //hid_t b=H5Pset_chunk (xfer_plist, 1, &buffsize);
  //hid_t b=H5Pset_buffer(xfer_plist, (hsize_t)buffsize, NULL, NULL);
  //unsigned szip_options_mask = H5_SZIP_NN_OPTION_MASK;
  //unsigned szip_pixels_per_block = 16;
  //hid_t b = H5Pset_szip (xfer_plist, szip_options_mask, szip_pixels_per_block);

  hid_t sid1  = H5Screate_simple(D,view.gcount,NULL);
  hid_t memspace=H5Screate_simple(D,view.count,NULL);

  //hsize_t blocks[D];
  //for(int i=0;i<D; ++i) blocks[i]=1;
  Timer myclock;
  herr_t ret;
  for(int i=0; i<nc; ++i)
  {
    sprintf(fname,"u_%d",i);
    hid_t dset_id=H5Dcreate(h_file,fname,etype,sid1,H5P_DEFAULT);
    hid_t filespace=H5Dget_space(dset_id);
    //if(chunked)
    //  ret=H5Sselect_hyperslab(filespace,H5S_SELECT_SET,view.offset,NULL,view.count,blocks);
    //else
    ret=H5Sselect_hyperslab(filespace,H5S_SELECT_SET,view.offset,NULL,view.count,NULL);
    ret = H5Dwrite(dset_id,etype,memspace,filespace,xfer_plist,first+i*view.size());
    H5Sclose(filespace);
    H5Dclose(dset_id);
  }
  app_log() << "Writing alone=" << myclock.elapsed() << endl;

  H5Sclose(sid1);
  H5Sclose(memspace);
  H5Pclose(xfer_plist);
  H5Pclose(h_plist);
  H5Fclose(h_file);
}

int main(int argc, char** argv)
{
  mpi::environment env(argc,argv);
  mpi::communicator mycomm;
  OhmmsInfo ohmms("dnsio",mycomm.rank(),0,1);

  ///number of iterations
  int niters=10;
  ///rows
  int sizeN=1024;
  ///partition x
  int nx=8;
  ///if true, debug the implemnetation
  bool debug=false;

  int ic=0;
  while(ic<argc)
  {
    std::string a(argv[ic]);
    if(a == "opt_s") 
      sizeN=atoi(argv[++ic]);
    else if(a == "opt_p") 
      nx=atoi(argv[++ic]);
    else if(a == "opt_i")
      niters=atoi(argv[++ic]);
    else if(a == "opt_debug")
      debug=true;
    ++ic;
  }

  int ny=mycomm.size()/nx;
  //nx=mycomm.size();
  //int ny=1;
  int locN=sizeN/nx;
  int locM=sizeN/ny;

  const bool iscomplex=false;
  const int nvars=1;
  typedef float value_type;

  //create my view
  DataView<3> myview(mycomm);
  myview.init(sizeN,nx,ny,iscomplex);
  //myview.init(sizeN,iscomplex);

  int ng_data=locN*locM*sizeN;
  //int ng_data=locN*sizeN;
  if(iscomplex) ng_data *=2;

  Matrix<value_type> A(nvars,ng_data);
  app_log() << "BW sizeN= " << sizeN << " px= " << nx << " py= " << ny << endl;
  app_log() << "BW LocN= " << locN << " LocM= " << locM << " localsize= " << A.size() <<endl;
  A=mycomm.rank();
  //for(int i=0; i<A.size(); ++i) A(i)=Random();

  double nbytes=static_cast<double>(A.size()*mycomm.size()*sizeof(value_type))/1024./1024./1024.;
  hsize_t buffsize=sizeN;

  Timer myclock;
  write_pencils_h5(myview,A.data(),nvars,"cube.0",false,true);
  //write_pencils_h5(A.cols(),nvars,buffsize,A.data(),"fun",mycomm);
  double dt=myclock.elapsed();
  app_log() << "BW col+nobuff GBtype=" << nbytes <<  " time= " << dt << " rate=" << nbytes/dt << " GBytes/sec" << endl;

  //myclock.restart();
  //write_pencils_h5(myview,A.data(),nvars,"cube.1",true,true);
  //dt=myclock.elapsed();
  //app_log() << "BW col+buff   GBtype=" << nbytes <<  " time= " << dt << " rate=" << nbytes/dt << " GBytes/sec" << endl;

  //char fname[32];
  //myclock.restart();
  //for(int i=0; i<nvars; ++i)
  //{
  //  sprintf(fname,"cube.v%i",i);
  //  write_pencils_h5(myview,A[i],1,fname,true,true);
  //}
  //dt=myclock.elapsed();
  //app_log() << "BW col+buff(6) GBtype=" << nbytes <<  " time= " << dt << " rate=" << nbytes/dt << " GBytes/sec" << endl;

  //myclock.restart();
  //write_pencils_h5(A.cols(),nvars,A.data(),"fun",mycomm);
  //dt=myclock.elapsed();

  //myclock.restart();
  //write_pencils_h5(A.size(),1,A.data(),"bigfun",mycomm);
  //dt1=myclock.elapsed();
  //app_log() << "BW H5 GBtype=" << nbytes <<  " time= " << dt << " rate=" << nbytes/dt << " GBytes/sec" << endl;
  //app_log() << "BW H5 GBtype=" << nbytes <<  " time= " << dt1 << " rate=" << nbytes/dt1 << " GBytes/sec" << endl;

  return 0;
}

